/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
#pragma once
#include <baja/graphics/dx11.1/details/common.hpp>
#include <dependencies/includes/dxtex.h>
#pragma comment(lib, "dxtex")

namespace baja { namespace graphics { namespace dx11_1 { namespace details {

inline WRL::ComPtr<ID3D11Resource> createTextureFromSystemMemory(
    ID3D11Device* d3d11Device,
    const DXGI_FORMAT format,
    const mipmap_load mipmapLoading,
    const D3D11_USAGE usage,
    const uint32 bindFlags,
    const uint32 cpuAccessFlags,
    const uint32 miscFlags,
    const DirectX::ScratchImage& image
    )
{
    BAJA_CHECK_ARG(format != DXGI_FORMAT_UNKNOWN);
    BAJA_CHECK_ARG(d3d11Device);

    DirectX::TexMetadata metadata = image.GetMetadata();
    const DirectX::Image* images = image.GetImages();
    size_t numImages = image.GetImageCount();
    BAJA_CHECK_ARG(metadata.mipLevels > 0 && metadata.arraySize > 0);
    BAJA_CHECK_ARG(images);
    BAJA_CHECK_ARG(numImages > 0);
    
    win32_result wr;

    DirectX::ScratchImage convertedImage;
    if (format != DXGI_FORMAT_FROM_FILE && format != metadata.format)
    {
        DirectX::convertPixelFormat(format, image, &convertedImage);
        metadata = convertedImage.GetMetadata();
        images = convertedImage.GetImages();
        numImages = convertedImage.GetImageCount();
    }

    DirectX::ScratchImage autoGeneratedMips;
    if (mipmapLoading == mipmap_load::noMipmaps)
    {
        // only load the top level image
        metadata.mipLevels = 1;
    }
    else if (mipmapLoading == mipmap_load::forceMipmaps)
    {
        DirectX::generateFullMipChain(images[0], &autoGeneratedMips);
        metadata = autoGeneratedMips.GetMetadata();
        images = autoGeneratedMips.GetImages();
        numImages = autoGeneratedMips.GetImageCount();
    }
    
    // fill out subresource array
    std::vector<D3D11_SUBRESOURCE_DATA> initData(metadata.mipLevels * metadata.arraySize);

    if (metadata.dimension == DirectX::TEX_DIMENSION_TEXTURE3D)
    {
        wr = (metadata.depth > 0);
        // Direct3D 11 doesn't support arrays of 3D textures
        if (metadata.arraySize > 1)
            wr = E_FAIL; // ERROR:
            
        size_t depth = metadata.depth;
        uint32 idx = 0;
        for (uint32 level = 0; level < metadata.mipLevels; level++)
        {
            size_t index = metadata.ComputeIndex(level, 0, 0);
            wr = (index < numImages);
                    
            const DirectX::Image& img = images[index];
            if (img.format != metadata.format)
            {
                wr = E_FAIL;
            }
            else if (img.pixels == NULL)
            {
                wr = E_POINTER;
            }
            else
            {
                // Verify pixels in image 1 .. (depth-1) are exactly image->slicePitch apart
                // For 3D textures, this relies on all slices of the same miplevel being continous in memory
                // (this is how ScratchImage lays them out), which is why we just give the 0th slice to Direct3D 11
                const uint8_t* pSlice = img.pixels + img.slicePitch;
                for (uint32 slice = 1; slice < depth; slice++)
                {
                    size_t tindex = metadata.ComputeIndex(level, 0, slice);
                    wr = (tindex < numImages);
                            
                    const DirectX::Image& timg = images[tindex];
                    if (timg.pixels == NULL)
                    {
                        wr = E_POINTER;
                    }
                    else if (timg.pixels != pSlice || timg.format != metadata.format || timg.rowPitch != img.rowPitch || timg.slicePitch != img.slicePitch)
                    {
                        wr = E_FAIL;
                    }
                    else 
                    {
                        pSlice = timg.pixels + img.slicePitch;
                    }
                }

                BAJA_VERIFY(idx < (metadata.mipLevels * metadata.arraySize));
                initData[idx].pSysMem = img.pixels;
                initData[idx].SysMemPitch = static_cast<DWORD>(img.rowPitch);
                initData[idx].SysMemSlicePitch = static_cast<DWORD>(img.slicePitch);
                ++idx;
                    
                if (depth > 1)
                {
                    depth >>= 1;
                }
            }
        }
    }
    else
    {
        //--- 1D or 2D texture case ---------------------------------------------------
            
        uint32 idx = 0;
        for (uint32 item = 0; item < metadata.arraySize; item++)
        {
            for (uint32 level = 0; level < metadata.mipLevels; level++)
            {
                size_t index = metadata.ComputeIndex(level, item, 0);
                wr = (index < numImages);
                const DirectX::Image& img = images[index];
                if (img.format != metadata.format)
                {
                    wr = E_FAIL;
                }
                else if (img.pixels == NULL)
                {
                    wr = E_POINTER;
                }
                else
                {
                    BAJA_VERIFY(idx < (metadata.mipLevels * metadata.arraySize));
                    initData[idx].pSysMem = img.pixels;
                    initData[idx].SysMemPitch = static_cast<DWORD>( img.rowPitch );
                    initData[idx].SysMemSlicePitch = static_cast<DWORD>( img.slicePitch );
                    ++idx;
                }
            }
        }
    }
    
    WRL::ComPtr<ID3D11Resource> d3d11Resource;

    switch (metadata.dimension)
    {
    case DirectX::TEX_DIMENSION_TEXTURE1D:
        {
            D3D11_TEXTURE1D_DESC desc;
            desc.Width = static_cast<UINT>(metadata.width);
            desc.MipLevels = static_cast<UINT>(metadata.mipLevels);
            desc.ArraySize = static_cast<UINT>(metadata.arraySize);
            desc.Format = metadata.format;
            desc.Usage = usage;
            desc.BindFlags = bindFlags;
            desc.CPUAccessFlags = cpuAccessFlags;
            desc.MiscFlags = miscFlags;
            
            wr = d3d11Device->CreateTexture1D(&desc, initData.data(), reinterpret_cast<ID3D11Texture1D**>(d3d11Resource.GetAddressOf()));
        }
        break;
    
    case DirectX::TEX_DIMENSION_TEXTURE2D:
        {
            D3D11_TEXTURE2D_DESC desc;
            desc.Width = static_cast<UINT>(metadata.width);
            desc.Height = static_cast<UINT>(metadata.height);
            desc.MipLevels = static_cast<UINT>(metadata.mipLevels);
            desc.ArraySize = static_cast<UINT>(metadata.arraySize);
            desc.Format = metadata.format;
            desc.SampleDesc.Count = 1;
            desc.SampleDesc.Quality = 0;
            desc.Usage = usage;
            desc.BindFlags = bindFlags;
            desc.CPUAccessFlags = cpuAccessFlags;

            desc.MiscFlags = miscFlags;
            desc.MiscFlags |= (metadata.miscFlags & DirectX::TEX_MISC_TEXTURECUBE) ? D3D11_RESOURCE_MISC_TEXTURECUBE : 0;

            wr = d3d11Device->CreateTexture2D(&desc, initData.data(), reinterpret_cast<ID3D11Texture2D**>(d3d11Resource.GetAddressOf()));
        }
        break;
    
    case DirectX::TEX_DIMENSION_TEXTURE3D:
        {
            D3D11_TEXTURE3D_DESC desc;
            desc.Width = static_cast<UINT>(metadata.width);
            desc.Height = static_cast<UINT>(metadata.height);
            desc.Depth = static_cast<UINT>(metadata.depth);
            desc.MipLevels = static_cast<UINT>(metadata.mipLevels);
            desc.Format = metadata.format;
            desc.Usage = usage;
            desc.BindFlags = bindFlags;
            desc.CPUAccessFlags = cpuAccessFlags;
            desc.MiscFlags = 0;

            wr = d3d11Device->CreateTexture3D(&desc, initData.data(), reinterpret_cast<ID3D11Texture3D**>(d3d11Resource.GetAddressOf()));
        }
        break;

    default:
        wr = E_FAIL;
        break;
    }

    return d3d11Resource;
}

inline uint32 computeScanlines(
    const DXGI_FORMAT format, 
    const uint32 height
    )
{
    switch (format)
    {
    case DXGI_FORMAT_BC1_TYPELESS:
    case DXGI_FORMAT_BC1_UNORM:
    case DXGI_FORMAT_BC1_UNORM_SRGB:
    case DXGI_FORMAT_BC2_TYPELESS:
    case DXGI_FORMAT_BC2_UNORM:
    case DXGI_FORMAT_BC2_UNORM_SRGB:
    case DXGI_FORMAT_BC3_TYPELESS:
    case DXGI_FORMAT_BC3_UNORM:
    case DXGI_FORMAT_BC3_UNORM_SRGB:
    case DXGI_FORMAT_BC4_TYPELESS:
    case DXGI_FORMAT_BC4_UNORM:
    case DXGI_FORMAT_BC4_SNORM:
    case DXGI_FORMAT_BC5_TYPELESS:
    case DXGI_FORMAT_BC5_UNORM:
    case DXGI_FORMAT_BC5_SNORM:
    case DXGI_FORMAT_BC6H_TYPELESS:
    case DXGI_FORMAT_BC6H_UF16:
    case DXGI_FORMAT_BC6H_SF16:
    case DXGI_FORMAT_BC7_TYPELESS:
    case DXGI_FORMAT_BC7_UNORM:
    case DXGI_FORMAT_BC7_UNORM_SRGB:
        return max(1, (height + 3) / 4);

    default:
        return height;
    }
}

inline static void _capture(
    ID3D11DeviceContext* d3d11Context,
    ID3D11Resource* d3d11Resource,
    const DirectX::TexMetadata& metadata,
    DirectX::ScratchImage* image
    )
{
    BAJA_CHECK_ARG(d3d11Context);
    BAJA_CHECK_ARG(d3d11Resource);
    BAJA_CHECK_ARG(image);
    
    win32_result wr;

    if (metadata.dimension == DirectX::TEX_DIMENSION_TEXTURE3D)
    {
        BAJA_VERIFY(metadata.arraySize == 1);
        
        size_t height = metadata.height;
        size_t depth = metadata.depth;
        for (size_t level = 0; level < metadata.mipLevels; level++)
        {
            D3D11_MAPPED_SUBRESOURCE mapped;
            uint32 dindex = D3D11CalcSubresource(static_cast<UINT>(level), static_cast<UINT>(0), static_cast<UINT>(metadata.mipLevels));
            const uint8_t* pslice = NULL;
            size_t lines = 0;
            bool isMapped = false;

            wr = d3d11Context->Map(d3d11Resource, dindex, D3D11_MAP_READ, 0, &mapped);
            isMapped = true;

            pslice = reinterpret_cast<const uint8_t*>(mapped.pData);
            if (!pslice) wr = E_FAIL; // ERROR:
            lines = static_cast<size_t>(computeScanlines(metadata.format, static_cast<baja::uint32>(height)));

            for (uint32 slice = 0; slice < depth; slice++)
            {
                const DirectX::Image* img = image->GetImage(level, 0, slice);
                if (!img) wr = E_FAIL; // ERROR: 
                if (!img->pixels) wr = E_FAIL; // ERROR: 

                const uint8_t* sptr = pslice;
                uint8_t* dptr = img->pixels;
                for (uint32 h = 0; h < lines; h++)
                {
                    size_t msize = min(img->rowPitch, mapped.RowPitch);
                    memcpy_s(dptr, img->rowPitch, sptr, msize);
                    sptr += mapped.RowPitch;
                    dptr += img->rowPitch;
                }
                
                pslice += mapped.DepthPitch;
            }

            if (isMapped)
            {
                d3d11Context->Unmap(d3d11Resource, dindex);
            }
            
            if (height > 1)
            {
                height >>= 1;
            }
            if (depth > 1)
            {
                depth >>= 1;
            }
        }
    }
    else
    {
        // 1D or 2D texture
        BAJA_VERIFY(metadata.depth == 1);
        for (uint32 item = 0; item < metadata.arraySize; item++)
        {
            size_t height = metadata.height;
            for (size_t level = 0; level < metadata.mipLevels; level++)
            {
                uint32 lines = 0;
                uint32 dindex = D3D11CalcSubresource(static_cast<UINT>(level), static_cast<UINT>(item), static_cast<UINT>(metadata.mipLevels));
                D3D11_MAPPED_SUBRESOURCE mapped;
                bool isMapped = false;

                wr = d3d11Context->Map(d3d11Resource, dindex, D3D11_MAP_READ, 0, &mapped);
                isMapped = true;

                const DirectX::Image* img = image->GetImage(level, item, 0);
                if (!img) wr = E_FAIL; // ERROR:
                if (!img->pixels) wr = E_FAIL; // ERROR:
                lines = computeScanlines(metadata.format, static_cast<baja::uint32>(height));
                
                const uint8_t* sptr = reinterpret_cast<const uint8_t*>(mapped.pData);
                uint8_t* dptr = img->pixels;
                for (uint32 h = 0; h < lines; h++)
                {
                    size_t msize = min(img->rowPitch, mapped.RowPitch);
                    memcpy_s(dptr, img->rowPitch, sptr, msize);
                    sptr += mapped.RowPitch;
                    dptr += img->rowPitch;
                }

                if (isMapped)
                {
                    d3d11Context->Unmap(d3d11Resource, dindex);
                }
                
                if (height > 1)
                {
                    height >>= 1;
                }
            }
        }
    }
}

inline void copyTextureToSystemMemory(
    ID3D11Device* d3d11Device,
    ID3D11DeviceContext* d3d11Context,
    ID3D11Resource* d3d11Resource,
    DirectX::ScratchImage* image
    )
{
    BAJA_CHECK_ARG(d3d11Device);
    BAJA_CHECK_ARG(d3d11Context);
    BAJA_CHECK_ARG(d3d11Resource);
    BAJA_CHECK_ARG(image);

    D3D11_RESOURCE_DIMENSION resType = D3D11_RESOURCE_DIMENSION_UNKNOWN;
    d3d11Resource->GetType(&resType);

    win32_result wr;

    WRL::ComPtr<ID3D11Resource> res = d3d11Resource;

    switch (resType)
    {
    case D3D11_RESOURCE_DIMENSION_TEXTURE1D:
        {
            WRL::ComPtr<ID3D11Texture1D> texture;
            res.As(&texture); // TEST: does this throw if interface is not supported ?
            
            D3D11_TEXTURE1D_DESC desc;
            texture->GetDesc(&desc);
            
            desc.BindFlags = 0;
            desc.MiscFlags = 0;
            desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
            desc.Usage = D3D11_USAGE_STAGING;
            
            WRL::ComPtr<ID3D11Texture1D> stagingTexture;
            wr = d3d11Device->CreateTexture1D(&desc, 0, stagingTexture.GetAddressOf());
            d3d11Context->CopyResource(stagingTexture.Get(), res.Get());
            
            DirectX::TexMetadata mdata;
            mdata.width = desc.Width;
            mdata.height = mdata.depth = 1;
            mdata.arraySize = desc.ArraySize;
            mdata.mipLevels = desc.MipLevels;
            mdata.miscFlags = 0;
            mdata.format = desc.Format;
            mdata.dimension = DirectX::TEX_DIMENSION_TEXTURE1D;
            
            wr = image->Initialize(mdata);
             _capture(d3d11Context, stagingTexture.Get(), mdata, image);
        }
        break;
    
    case D3D11_RESOURCE_DIMENSION_TEXTURE2D:
        {
            WRL::ComPtr<ID3D11Texture2D> texture;
            res.As(&texture);

            D3D11_TEXTURE2D_DESC desc;
            texture->GetDesc(&desc);
            
            WRL::ComPtr<ID3D11Texture2D> stagingTexture;
            if (desc.SampleDesc.Count > 1)
            {
                desc.SampleDesc.Count = 1;
                desc.SampleDesc.Quality = 0;
                
                WRL::ComPtr<ID3D11Texture2D> tmp;
                wr = d3d11Device->CreateTexture2D(&desc, 0, tmp.GetAddressOf());
                for (uint32 item = 0; item < desc.ArraySize; item++)
                {
                    for (uint32 level = 0; level < desc.MipLevels; level++)
                    {
                        uint32 index = D3D11CalcSubresource(level, item, desc.MipLevels);
                        d3d11Context->ResolveSubresource(tmp.Get(), index, d3d11Resource, index, desc.Format);
                    }
                }
                
                desc.BindFlags = 0;
                desc.MiscFlags &= D3D11_RESOURCE_MISC_TEXTURECUBE;
                desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
                desc.Usage = D3D11_USAGE_STAGING;
                
                wr = d3d11Device->CreateTexture2D(&desc, 0, stagingTexture.GetAddressOf());
                d3d11Context->CopyResource(stagingTexture.Get(), tmp.Get());
            }
            else
            {
                desc.BindFlags = 0;
                desc.MiscFlags &= D3D11_RESOURCE_MISC_TEXTURECUBE;
                desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
                desc.Usage = D3D11_USAGE_STAGING;
                
                wr = d3d11Device->CreateTexture2D(&desc, 0, stagingTexture.GetAddressOf());
                d3d11Context->CopyResource(stagingTexture.Get(), res.Get());
            }

            DirectX::TexMetadata mdata;
            mdata.width = desc.Width;
            mdata.height = desc.Height;
            mdata.depth = 1;
            mdata.arraySize = desc.ArraySize;
            mdata.mipLevels = desc.MipLevels;
            mdata.miscFlags = (desc.MiscFlags & D3D11_RESOURCE_MISC_TEXTURECUBE) ? DirectX::TEX_MISC_TEXTURECUBE : 0;
            mdata.format = desc.Format;
            mdata.dimension = DirectX::TEX_DIMENSION_TEXTURE2D;
            
            wr = image->Initialize(mdata);
            _capture(d3d11Context, stagingTexture.Get(), mdata, image);
        }
        break;
    
    /*
    // NYI:
    case D3D11_RESOURCE_DIMENSION_TEXTURE3D:
        {
            WRL::ComPtr<ID3D11Texture3D> texture;
            d3d11Resource.As(&texture);
            
            D3D11_TEXTURE3D_DESC desc;
            texture->GetDesc(&desc);
            
            desc.BindFlags = 0;
            desc.MiscFlags = 0;
            desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
            desc.Usage = D3D11_USAGE_STAGING;
            
            WRL::ComPtr<ID3D11Texture3D> stagingTexture;
            wr = d3d11Device->CreateTexture3D(&desc, 0, stagingTexture.GetAddressOf());
            d3d11Context->CopyResource(stagingTexture, d3d11Resource.Get());
            TexMetadata mdata;
            mdata.width = desc.Width;
            mdata.height = desc.Height;
            mdata.depth = desc.Depth;
            mdata.arraySize = 1;
            mdata.mipLevels = desc.MipLevels;
            mdata.miscFlags = 0;
            mdata.format = desc.Format;
            mdata.dimension = TEX_DIMENSION_TEXTURE3D;

            wr = image->Initialize(mdata);
            _capture(d3d11Context, stagingTexture, mdata, image);
        }
        break;
    */

    default:
        wr = E_FAIL;
        break;
    }
}

}}}} // namespace baja::graphics::dx11_1::details
